import json
import numpy as np
import math
from scipy.ndimage import distance_transform_edt
from collections import Counter
from PIL import Image
from random import *
from utils import *
from config import *

def read_json(filename):
    with open(filename, 'r') as file:
        data = json.load(file)
    return data

def euler_to_rotation_matrix(euler):
    rx, ry, rz = euler
    Rx = np.array([
        [1, 0, 0],
        [0, math.cos(rx), -math.sin(rx)],
        [0, math.sin(rx), math.cos(rx)]
    ])
    Ry = np.array([
        [math.cos(ry), 0, math.sin(ry)],
        [0, 1, 0],
        [-math.sin(ry), 0, math.cos(ry)]
    ])
    Rz = np.array([
        [math.cos(rz), -math.sin(rz), 0],
        [math.sin(rz), math.cos(rz), 0],
        [0, 0, 1]
    ])
    R = Rz @ Ry @ Rx
    return R

def build_camera_matrix(location, rotation_euler):
    R = euler_to_rotation_matrix(rotation_euler)
    T = np.array(location).reshape(3,1)
    M = np.eye(4)
    M[:3, :3] = R
    M[:3, 3] = T.flatten()
    return M

def world_to_camera(loc_world, cam_location, cam_rotation):
    M = build_camera_matrix(cam_location, cam_rotation)
    M_inv = np.linalg.inv(M)
    point_world = np.array([loc_world[0], loc_world[1], loc_world[2], 1.0])

    point_cam = M_inv @ point_world
    adjust = np.array([
        [1, 0, 0],
        [0, -1, 0],
        [0,  0, -1]
    ])
    point_cam = adjust @ point_cam[:3]

    return point_cam  # 返回 x, y, z

def project_point(point_cam, fov, image_width, image_height):
    x, y, z = point_cam
    if z <= 0:
        return None
    fov_x, fov_y = fov
    fx = image_width / (2 * math.tan(fov_x/2))
    fy = image_height / (2 * math.tan(fov_y/2))
    cx = image_width / 2.0
    cy = image_height / 2.0
    u = fx * (x / z) + cx
    v = fy * (y / z) + cy
    if u < 0 or u >= image_width or v < 0 or v >= image_height:
        return None
    return (u, v)

def world_to_camera_pixel(loc_world, cam_info, image_width = 960, image_height = 540): # WIDTH, HEIGHT
    cam_location = cam_info.get("location")
    cam_rotation = cam_info.get("rotation")
    fov = cam_info.get("fov")
    
    point_cam = world_to_camera(loc_world, cam_location, cam_rotation)
    
    # point_cam = -point_cam

    pixel = project_point(point_cam, fov, image_width, image_height)
    return pixel

def mask_decoder(rle):
    counts = rle.get("counts", [])
    size = rle.get("size", [])

    height, width = size
    total_pixels = height * width

    flat_mask = np.zeros(total_pixels, dtype=np.uint8)

    current_index = 0
    current_val = rle.get("prefix", 0)
    for count in counts:
        if current_index + count > total_pixels:
            raise ValueError("RLE Error")
        flat_mask[current_index:current_index+count] = current_val
        current_index += count
        current_val = 1 - current_val

    if current_index != total_pixels:
        raise ValueError("RLE Error")

    mask = flat_mask.reshape((height, width))
    return mask

def get_mask_center(mask): # HEIGHT, WIDTH
    dist = distance_transform_edt(mask)
    center = np.unravel_index(np.argmax(dist), mask.shape)
    return (center[1], center[0])

def get_main_cate(objs):
    categories = [item['cate'] for item in objs]
    category_counts = Counter(categories)
    most_common_category, most_common_count = category_counts.most_common(1)[0]
    if most_common_count > 1:
        return most_common_category, most_common_count
    else:
        return None, None

def check_nomask(data):
    nomask_count = 0
    for id, info in data.items():
        if info['mask'].any() == False:
            nomask_count += 1
    return nomask_count <= 1

def check_depth(path):
    image = Image.open(path)
    image_array = np.array(image)
    last_row = image_array[-1, :, :]
    return not np.all(last_row == [255, 255, 255])

def check_enough_visible_part(info, cam):
    '''
    Check if the object has enough visible parts in the camera view (at least 4 points of bbox).
    '''
    xs = [info['bboxmin'][0], info['bboxmax'][0]]
    ys = [info['bboxmin'][1], info['bboxmax'][1]]
    zs = [info['bboxmin'][2], info['bboxmax'][2]]
    count = 0
    for x in xs:
        for y in ys:
            for z in zs:
                loc = (x, y, z)
                pixel = world_to_camera_pixel(loc, cam)
                if pixel is not None:
                    count += 1
    return count >= 2

def check_bottom_visible(info, cam):
    '''
    check if the bottom of the object is visible in the camera view (zs[0] only)
    '''
    xs = [info['bboxmin'][0], info['bboxmax'][0]]
    ys = [info['bboxmin'][1], info['bboxmax'][1]]
    count = 0
    for x in xs:
        for y in ys:
            loc = (x, y, info['bboxmin'][2])
            pixel = world_to_camera_pixel(loc, cam)
            if pixel is not None:
                count += 1
    return count == 4

def get_angle(x, y):
    return math.atan2(y, x)

def get_original_id(id):
    if '_' in id:
        return id.split('_')[0]
    return id

def calc_rank(base_objs, objs, subset, color):
    group_count = len(objs)
    info_list = []
    for oid, info in objs.items():
        x_min = info["bboxmin"][0]
        x_max = info["bboxmax"][0]
        y_min = info["bboxmin"][1]
        y_max = info["bboxmax"][1]
        height = info["bboxmax"][2] - info["bboxmin"][2]
        info_list.append((oid, x_min, x_max, y_min, y_max, height))
    
    for current in info_list:
        oid, x_min, x_max, y_min, y_max, h = current

        ambiguous_x = False
        ambiguous_y = False
        ambiguous_h = False

        for other in info_list:
            other_id, other_x_min, other_x_max, _, _, _ = other
            if oid == other_id:
                continue
            if not (x_max <= other_x_min or other_x_max <= x_min):
                ambiguous_x = True
                break

        for other in info_list:
            other_id, _, _, other_y_min, other_y_max, _ = other
            if oid == other_id:
                continue
            if not (y_max <= other_y_min or other_y_max <= y_min):
                ambiguous_y = True
                break
        
        for other in info_list:
            other_id, _, _, _, _, other_h = other
            if oid == other_id:
                continue
            if math.fabs(h - other_h) / h < 0.15:
                ambiguous_h = True
                break

        if not ambiguous_x:
            left_count = sum(1 for other in info_list if other[0] != oid and other[2] <= x_min)
            base_objs[oid]['exclusive_ref'].append({
                'color': color,
                'title': subset,
                'rank': left_count + 1,
                'all': group_count,
                'dir': 'x'
            })
        if group_count == 1:
            break

        if not ambiguous_y:
            lower_count = sum(1 for other in info_list if other[0] != oid and other[4] <= y_min)
            base_objs[oid]['exclusive_ref'].append({
                'color': color,
                'title': subset,
                'rank': lower_count + 1,
                'all': group_count,
                'dir': 'y'
            })
        
        if not ambiguous_h:
            lower_count = sum(1 for other in info_list if other[0] != oid and other[5] <= h)
            base_objs[oid]['exclusive_ref'].append({
                'color': color,
                'title': subset,
                'rank': lower_count + 1,
                'all': group_count,
                'dir': 'z'
            })

def get_rank(x):
    if x == 1:
        return 'first'
    elif x == 2:
        return 'second'
    elif x == 3:
        return 'third'
    else:
        return f'{x}-th'

def format_coordinatexy(x, y):
    return f'({x / WIDTH:.3f}, {y / HEIGHT:.3f})'
def format_coordinate(info):
    return f'({info["2dloc"][0] / WIDTH:.3f}, {info["2dloc"][1] / HEIGHT:.3f})'

def ref_to_name(ref):
    name = ref['title']
    if ref['color'] != '':
        name = f'{ref["color"]} {name}'
    return name

def ref_to_desc(ref):
    name = ref_to_name(ref)
    if ref['all'] == 1:
        return f'the {name}'
    if ref['dir'] == 'y':
        if ref['rank'] == 1:
            desc = f'the leftmost {name}' if ref['all'] > 2 else f'the left {name}'
        elif ref['rank'] == ref['all']:
            desc = f'the rightmost {name}' if ref['all'] > 2 else f'the right {name}'
        else:
            if random() < 0.5:
                rk = get_rank(ref['rank'])
                desc = f'the {rk} {name} from left to right'
            else:
                rk = get_rank(ref['all'] - ref['rank'] + 1)
                desc = f'the {rk} {name} from right to left'
    elif ref['dir'] == 'x':
        if ref['rank'] == 1:
            desc = f'the farthest {name}' if ref['all'] > 2 else f'the farther {name}'
        elif ref['rank'] == ref['all']:
            desc = f'the nearest {name}' if ref['all'] > 2 else f'the nearer {name}'
        else:
            if random() < 0.5:
                rk = get_rank(ref['rank'])
                desc = f'the {rk} {name} from far to near'
            else:
                rk = get_rank(ref['all'] - ref['rank'] + 1)
                desc = f'the {rk} {name} from near to far'
    else:
        adj1 = 'short'
        adj2 = 'tall'
        if ref['color'] != '':
            adj1 = 'small'
            adj2 = 'big'

        if ref['rank'] == 1:
            if adj1 == 'short':
                desc = f'the shortest {name}' if ref['all'] > 2 else f'the shorter {name}'
            else:
                desc = f'the smallest {name}' if ref['all'] > 2 else f'the smaller {name}'
        elif ref['rank'] == ref['all']:
            if adj1 == 'short':
                desc = f'the tallest {name}' if ref['all'] > 2 else f'the taller {name}'
            else:
                desc = f'the biggest {name}' if ref['all'] > 2 else f'the bigger {name}'
        else:
            if random() < 0.5:
                rk = get_rank(ref['rank'])
                desc = f'the {rk} {name} from {adj1} to {adj2}'
            else:
                rk = get_rank(ref['all'] - ref['rank'] + 1)
                desc = f'the {rk} {name} from {adj2} to {adj1}'
    return desc

def ref_to_thinking(obj_dict, info, ref, desc):
    name = ref_to_name(ref)
    if ref['all'] == 1:
        return f'There is only one {name} in the image, which is at {format_coordinate(info)}.'
    text = f'There is/are {ref["all"]} {name}(s) in the image. They are at: \n'
    for id, infoi in obj_dict.items():
        if not infoi['visible']:
            continue
        if ref['color'] != '' and infoi['color'] != ref['color']:
            continue
        if ref['title'] != 'object' and infoi['cate'] != ref['title']:
            continue
        text += f'pixel position at {format_coordinate(infoi)}'
        if ref['dir'] == 'z':
            text += f', with height {infoi["bboxmax"][2] - infoi["bboxmin"][2]:.3f}m'
        text += '\n'
    text += f'{desc} is the one at {format_coordinate(info)}.'
    return text

def float_eq(x, y):
    return math.fabs(x - y) < 1e-4

def check_table_pos(pos, cam, table):
    tabminx, tabminy, tabmaxx, tabmaxy, tablez = table
    if pos[0] < tabminx or pos[0] > tabmaxx or pos[1] < tabminy or pos[1] > tabmaxy:
        return None
    return world_to_camera_pixel((pos[0], pos[1], tablez), cam)

def pix_on_mask(pix, obj_dict):
    pix = (int(pix[0]), int(pix[1]))
    for id, info in obj_dict.items():
        if not info['visible']:
            continue
        if info['mask'][pix[1], pix[0]] == 1:
            return True
    return False

def calc_distance(obj1, obj2):
    bboxmin1 = obj1['bboxmin']
    bboxmax1 = obj1['bboxmax']
    bboxmin2 = obj2['bboxmin']
    bboxmax2 = obj2['bboxmax']
    dx = max(0, max(bboxmin1[0], bboxmin2[0]) - min(bboxmax1[0], bboxmax2[0]))
    dy = max(0, max(bboxmin1[1], bboxmin2[1]) - min(bboxmax1[1], bboxmax2[1]))
    distance = math.sqrt(dx**2 + dy**2)
    return distance
